package com.tzg157.fitness.modals;

import org.apache.spark.ml.Pipeline;
import org.apache.spark.ml.PipelineModel;
import org.apache.spark.ml.PipelineStage;
import org.apache.spark.ml.classification.RandomForestClassifier;
import org.apache.spark.ml.evaluation.RegressionEvaluator;
import org.apache.spark.ml.feature.StandardScaler;
import org.apache.spark.ml.feature.VectorAssembler;
import org.apache.spark.ml.regression.RandomForestRegressor;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.RowFactory;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;

public class RandomForestModal {
    private static final Logger logger = LoggerFactory.getLogger(RandomForestModal.class);

    public static void main(String[] args) throws IOException {
        System.setProperty("HADOOP_USER_NAME","root");
        SparkSession spark = SparkSession.builder()
                .appName("FitnessModel")
                .master("local[*]")
                .enableHiveSupport()
                .getOrCreate();
        // 创建SparkSession
        Dataset<Row> data = spark.sql("SELECT Age as age, Gender as gender, Weight as weight, Height as height," +
                " Max_BPM as max_bpm, Avg_BPM as avg_bpm, Resting_BPM as resting_bpm, Session_Duration as session_duration," +
                " Fat_Percentage as fat_percentage, Water_Intake as water_intake, Workout_Frequency as workout_frequency," +
                " Experience_Level as experience_level, BMI as bmi, Calories_Burned as calories_burned, Workout_Type as workout_type " +
                " FROM default.fitness_data");

        // 定义特征列名
        String[] featureColumns = new String[]{
                "age", "gender", "weight", "height", "max_bpm", "avg_bpm", "resting_bpm",
                "session_duration", "fat_percentage", "water_intake", "workout_frequency",
                "experience_level", "bmi"
        };

        // 检查并处理空值
        Dataset<Row> cleanedData = data.na().drop();

        // 组合特征列
        VectorAssembler assembler = new VectorAssembler()
                .setInputCols(featureColumns)
                .setOutputCol("features").setHandleInvalid("skip");

        // 标准化特征
        StandardScaler scaler = new StandardScaler()
                .setInputCol("features")
                .setOutputCol("scaledFeatures")
                .setWithStd(true)
                .setWithMean(true);

        // 训练卡路里消耗预测模型
        RandomForestRegressor rfCalories = new RandomForestRegressor()
                .setLabelCol("calories_burned")
                .setFeaturesCol("scaledFeatures")
                .setNumTrees(30)
                .setMaxDepth(4)
                .setSeed(12345);

        // 训练每次训练时长预测模型
        RandomForestRegressor rfDuration = new RandomForestRegressor()
                .setLabelCol("session_duration")
                .setFeaturesCol("scaledFeatures")
                .setNumTrees(30)
                .setMaxDepth(4)
                .setSeed(12345);

        // 训练每周训练次数预测模型
        RandomForestRegressor rfFrequency = new RandomForestRegressor()
                .setLabelCol("workout_frequency")
                .setFeaturesCol("scaledFeatures")
                .setNumTrees(30)
                .setMaxDepth(4)
                .setSeed(12345);

        // 训练锻炼期间心率预测模型
        RandomForestRegressor rfHeartRate = new RandomForestRegressor()
                .setLabelCol("max_bpm")
                .setFeaturesCol("scaledFeatures")
                .setNumTrees(30)
                .setMaxDepth(4)
                .setSeed(12345);

        // 训练健身方法推荐模型
        RandomForestClassifier rfWorkoutType = new RandomForestClassifier()
                .setLabelCol("workout_type")
                .setFeaturesCol("scaledFeatures")
                .setNumTrees(30)
                .setMaxDepth(4)
                .setSeed(12345);

        // 构建Pipeline
        Pipeline pipelineCalories = new Pipeline().setStages(new PipelineStage[]{assembler, scaler, rfCalories});
        Pipeline pipelineDuration = new Pipeline().setStages(new PipelineStage[]{assembler, scaler, rfDuration});
        Pipeline pipelineFrequency = new Pipeline().setStages(new PipelineStage[]{assembler, scaler, rfFrequency});
        Pipeline pipelineHeartRate = new Pipeline().setStages(new PipelineStage[]{assembler, scaler, rfHeartRate});
        Pipeline pipelineWorkoutType = new Pipeline().setStages(new PipelineStage[]{assembler, scaler, rfWorkoutType});

        // 分割数据集
        Dataset<Row>[] splits = cleanedData.randomSplit(new double[]{0.8, 0.2});
        Dataset<Row> trainingData = splits[0];
        Dataset<Row> testData = splits[1];

        // 训练模型
        PipelineModel modelCalories = pipelineCalories.fit(trainingData);
        modelCalories.write().overwrite().save("/opt/hive/models/random_forest/calories");
        PipelineModel modelDuration = pipelineDuration.fit(trainingData);
        modelDuration.write().overwrite().save("/opt/hive/models/random_forest/duration");
        PipelineModel modelFrequency = pipelineFrequency.fit(trainingData);
        modelFrequency.write().overwrite().save("/opt/hive/models/random_forest/frequency");
        PipelineModel modelHeartRate = pipelineHeartRate.fit(trainingData);
        modelHeartRate.write().overwrite().save("/opt/hive/models/random_forest/heartRate");
        PipelineModel modelWorkoutType = pipelineWorkoutType.fit(trainingData);
        modelWorkoutType.write().overwrite().save("/opt/hive/models/random_forest/workoutType");

        // 使用 testData 进行评估
        Dataset<Row> testPredictionsCalories = modelCalories.transform(testData);
        Dataset<Row> testPredictionsDuration = modelDuration.transform(testData);
        Dataset<Row> testPredictionsFrequency = modelFrequency.transform(testData);
        Dataset<Row> testPredictionsHeartRate = modelHeartRate.transform(testData);
        Dataset<Row> testPredictionsWorkoutType = modelWorkoutType.transform(testData);

        // 创建 RegressionEvaluator 实例
        RegressionEvaluator evaluator = new RegressionEvaluator()
                .setLabelCol("calories_burned")
                .setPredictionCol("prediction")
                .setMetricName("mse");

        // 计算均方误差 (MSE)
        double mseCalories = evaluator.evaluate(testPredictionsCalories);
        System.out.println("卡路里消耗预测的均方误差: " + mseCalories);

        // 重新设置 evaluator 的标签列以评估其他模型
        evaluator.setLabelCol("session_duration");
        double mseDuration = evaluator.evaluate(testPredictionsDuration);
        System.out.println("每次训练时长预测的均方误差: " + mseDuration);

        evaluator.setLabelCol("workout_frequency");
        double mseFrequency = evaluator.evaluate(testPredictionsFrequency);
        System.out.println("每周训练次数预测的均方误差: " + mseFrequency);

        evaluator.setLabelCol("max_bpm");
        double mseHeartRate = evaluator.evaluate(testPredictionsHeartRate);
        System.out.println("锻炼期间心率预测的均方误差: " + mseHeartRate);
    }
}
